/*
 * Copyright 2021 Google LLC.
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     https://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

/**
 *  Plotting of decision trees generated by TF-DF.
 *
 *  A tree is a recursive structure of node objects.
 *  A node contains one or more of the following components:
 *
 *    - A value: Representing the output of the node. If the node is not a leaf,
 *      the value is only present for analysis i.e. it is not used for
 *      predictions.
 *
 *    - A condition : For non-leaf nodes, the condition (also known as split)
 *      defines a binary test to branch to the positive or negative child.
 *
 *    - An explanation: Generally a plot showing the relation between the label
 *      and the condition to give insights about the effect of the condition.
 *
 *    - Two children : For non-leaf nodes, the children nodes. The first
 *      children (i.e. "node.children[0]") is the negative children (drawn in
 *      red). The second children is the positive one (drawn in green).
 *
 */

/**
 * Plots a single decision tree into a DOM element.
 * @param {!options} options Dictionary of configurations.
 * @param {!tree} raw_tree Recursive tree structure.
 * @param {string} canvas_id Id of the output dom element.
 */
function display_tree(options, raw_tree, canvas_id) {
  console.log(options);

  // Determine the node placement.
  const tree_struct = d3.tree().nodeSize(
      [options.node_y_offset, options.node_x_offset])(d3.hierarchy(raw_tree));

  // Boundaries of the node placement.
  let x_min = Infinity;
  let x_max = -x_min;
  let y_min = Infinity;
  let y_max = -x_min;

  tree_struct.each(d => {
    if (d.x > x_max) x_max = d.x;
    if (d.x < x_min) x_min = d.x;
    if (d.y > y_max) y_max = d.y;
    if (d.y < y_min) y_min = d.y;
  });

  // Size of the plot.
  const width = y_max - y_min + options.node_x_size + options.margin * 2;
  const height = x_max - x_min + options.node_y_size + options.margin * 2 +
      options.node_y_offset - options.node_y_size;

  const plot = d3.select(canvas_id);

  // Tool tip
  options.tooltip = plot.append('div')
                        .attr('width', 100)
                        .attr('height', 100)
                        .style('padding', '4px')
                        .style('background', '#fff')
                        .style('box-shadow', '4px 4px 0px rgba(0,0,0,0.1)')
                        .style('border', '1px solid black')
                        .style('font-family', 'sans-serif')
                        .style('font-size', options.font_size)
                        .style('position', 'absolute')
                        .style('z-index', '10')
                        .attr('pointer-events', 'none')
                        .style('display', 'none');

  // Create canvas
  const svg = plot.append('svg').attr('width', width).attr('height', height);
  const graph =
      svg.style('overflow', 'visible')
          .append('g')
          .attr('font-family', 'sans-serif')
          .attr('font-size', options.font_size)
          .attr(
              'transform',
              () => `translate(${options.margin},${
                  - x_min + options.node_y_offset / 2 + options.margin})`);

  // Plot bounding box.
  if (options.show_plot_bounding_box) {
    svg.append('rect')
        .attr('width', width)
        .attr('height', height)
        .attr('fill', 'none')
        .attr('stroke-width', 1.0)
        .attr('stroke', 'black');
  }

  // Draw the edges.
  display_edges(options, graph, tree_struct);

  // Draw the nodes.
  display_nodes(options, graph, tree_struct);
}

/**
 * Draw the nodes of the tree.
 * @param {!options} options Dictionary of configurations.
 * @param {!graph} graph D3 search handle containing the graph.
 * @param {!tree_struct} tree_struct Structure of the tree (node placement,
 *     data, etc.).
 */
function display_nodes(options, graph, tree_struct) {
  const nodes = graph.append('g')
                    .selectAll('g')
                    .data(tree_struct.descendants())
                    .join('g')
                    .attr('transform', d => `translate(${d.y},${d.x})`);

  nodes.append('rect')
      .attr('x', 0.5)
      .attr('y', 0.5)
      .attr('width', options.node_x_size)
      .attr('height', options.node_y_size)
      .attr('stroke', 'lightgrey')
      .attr('stroke-width', 1)
      .attr('fill', 'white')
      .attr('y', -options.node_y_size / 2);

  // Brackets on the right of condition nodes without children.
  non_leaf_node_without_children =
      nodes.filter(node => node.data.condition != null && node.children == null)
          .append('g')
          .attr('transform', `translate(${options.node_x_size},0)`);

  non_leaf_node_without_children.append('path')
      .attr('d', 'M0,0 C 10,0 0,10 10,10')
      .attr('fill', 'none')
      .attr('stroke-width', 1.0)
      .attr('stroke', '#F00');

  non_leaf_node_without_children.append('path')
      .attr('d', 'M0,0 C 10,0 0,-10 10,-10')
      .attr('fill', 'none')
      .attr('stroke-width', 1.0)
      .attr('stroke', '#0F0');

  const node_content = nodes.append('g').attr(
      'transform',
      `translate(0,${options.node_padding - options.node_y_size / 2})`);

  node_content.append(node => create_node_element(options, node));
}

/**
 * Creates the D3 content for a single node.
 * @param {!options} options Dictionary of configurations.
 * @param {!node} node Node to draw.
 * @return {!d3} D3 content.
 */
function create_node_element(options, node) {
  // Output accumulator.
  let output = {
    // Content to draw.
    content: d3.create('svg:g'),
    // Vertical offset to the next element to draw.
    vertical_offset: 0
  };

  // Conditions.
  if (node.data.condition != null) {
    display_condition(options, node.data.condition, output);
  }

  // Values.
  if (node.data.value != null) {
    display_value(options, node.data.value, output);
  }

  // Explanations.
  if (node.data.explanation != null) {
    display_explanation(options, node.data.explanation, output);
  }

  return output.content.node();
}


/**
 * Adds a single line of text inside of a node.
 * @param {!options} options Dictionary of configurations.
 * @param {string} text Text to display.
 * @param {!output} output Output display accumulator.
 */
function display_node_text(options, text, output) {
  output.content.append('text')
      .attr('x', options.node_padding)
      .attr('y', output.vertical_offset)
      .attr('alignment-baseline', 'hanging')
      .text(text);
  output.vertical_offset += 10;
}

/**
 * Adds a single line of text inside of a node with a tooltip.
 * @param {!options} options Dictionary of configurations.
 * @param {string} text Text to display.
 * @param {string} tooltip Text in the Tooltip.
 * @param {!output} output Output display accumulator.
 */
function display_node_text_with_tooltip(options, text, tooltip, output) {
  const item = output.content.append('text')
                   .attr('x', options.node_padding)
                   .attr('alignment-baseline', 'hanging')
                   .text(text);

  add_tooltip(options, item, () => tooltip);
  output.vertical_offset += 10;
}

/**
 * Adds a tooltip to a dom element.
 * @param {!options} options Dictionary of configurations.
 * @param {!dom} target Dom element to equip with a tooltip.
 * @param {!func} get_content Generates the html content of the tooltip.
 */
function add_tooltip(options, target, get_content) {
  function show(d) {
    options.tooltip.style('display', 'block');
    options.tooltip.html(get_content());
  }

  function hide(d) {
    options.tooltip.style('display', 'none');
  }

  function move(d) {
    options.tooltip.style('display', 'block');
    options.tooltip.style('left', (d.pageX + 5) + 'px');
    options.tooltip.style('top', d.pageY + 'px');
  }

  target.on('mouseover', show);
  target.on('mouseout', hide);
  target.on('mousemove', move);
}

/**
 * Adds a condition inside of a node.
 * @param {!options} options Dictionary of configurations.
 * @param {!condition} condition Condition to display.
 * @param {!output} output Output display accumulator.
 */
function display_condition(options, condition, output) {
  threshold_format = d3.format('r');

  if (condition.type === 'IS_MISSING') {
    display_node_text(options, `${condition.attribute} is missing`, output);
    return;
  }

  if (condition.type === 'IS_TRUE') {
    display_node_text(options, `${condition.attribute} is true`, output);
    return;
  }

  if (condition.type === 'NUMERICAL_IS_HIGHER_THAN') {
    format = d3.format('r');
    display_node_text(
        options,
        `${condition.attribute} >= ${threshold_format(condition.threshold)}`,
        output);
    return;
  }

  if (condition.type === 'CATEGORICAL_IS_IN') {
    display_node_text_with_tooltip(
        options, `${condition.attribute} in [...]`,
        `${condition.attribute} in [${condition.mask}]`, output);
    return;
  }

  if (condition.type === 'CATEGORICAL_SET_CONTAINS') {
    display_node_text_with_tooltip(
        options, `${condition.attribute} intersect [...]`,
        `${condition.attribute} intersect [${condition.mask}]`, output);
    return;
  }

  if (condition.type === 'NUMERICAL_SPARSE_OBLIQUE') {
    display_node_text_with_tooltip(
        options, `Sparse oblique split...`,
        `[${condition.attributes}]*[${condition.weights}]>=${
            threshold_format(condition.threshold)}`,
        output);
    return;
  }

  display_node_text(
      options, `Non supported condition ${condition.type}`, output);
}

/**
 * Adds a value inside of a node.
 * @param {!options} options Dictionary of configurations.
 * @param {!value} value Value to display.
 * @param {!output} output Output display accumulator.
 */
function display_value(options, value, output) {
  if (value.type === 'PROBABILITY') {
    const left_margin = 0;
    const right_margin = 50;
    const plot_width = options.node_x_size - options.node_padding * 2 -
        left_margin - right_margin;

    let cusum = Array.from(d3.cumsum(value.distribution));
    cusum.unshift(0);
    const distribution_plot = output.content.append('g').attr(
        'transform', `translate(0,${output.vertical_offset + 0.5})`);

    distribution_plot.selectAll('rect')
        .data(value.distribution)
        .join('rect')
        .attr('height', 10)
        .attr(
            'x',
            (d, i) =>
                (cusum[i] * plot_width + left_margin + options.node_padding))
        .attr('width', (d, i) => d * plot_width)
        .style('fill', (d, i) => d3.schemeSet1[i]);

    const num_examples =
        output.content.append('g')
            .attr('transform', `translate(0,${output.vertical_offset})`)
            .append('text')
            .attr('x', options.node_x_size - options.node_padding)
            .attr('alignment-baseline', 'hanging')
            .attr('text-anchor', 'end')
            .text(`(${value.num_examples})`);

    const distribution_details = d3.create('ul');
    distribution_details.selectAll('li')
        .data(value.distribution)
        .join('li')
        .append('span')
        .text(
            (d, i) =>
                'class ' + i + ': ' + d3.format('.3%')(value.distribution[i]));

    add_tooltip(options, distribution_plot, () => distribution_details.html());
    add_tooltip(options, num_examples, () => 'Number of examples');

    output.vertical_offset += 10;
    return;
  }

  if (value.type === 'REGRESSION') {
    display_node_text(
        options,
        'value: ' + d3.format('r')(value.value) + ` (` +
            d3.format('.6')(value.num_examples) + `)`,
        output);
    return;
  }

  display_node_text(options, `Non supported value ${value.type}`, output);
}

/**
 * Adds an explanation inside of a node.
 * @param {!options} options Dictionary of configurations.
 * @param {!explanation} explanation Explanation to display.
 * @param {!output} output Output display accumulator.
 */
function display_explanation(options, explanation, output) {
  // Margin before the explanation.
  output.vertical_offset += 10;

  display_node_text(
      options, `Non supported explanation ${explanation.type}`, output);
}


/**
 * Draw the edges of the tree.
 * @param {!options} options Dictionary of configurations.
 * @param {!graph} graph D3 search handle containing the graph.
 * @param {!tree_struct} tree_struct Structure of the tree (node placement,
 *     data, etc.).
 */
function display_edges(options, graph, tree_struct) {
  // Draw an edge between a parent and a child node with a bezier.
  function draw_single_edge(d) {
    return 'M' + (d.source.y + options.node_x_size) + ',' + d.source.x + ' C' +
        (d.source.y + options.node_x_size + options.edge_rounding) + ',' +
        d.source.x + ' ' + (d.target.y - options.edge_rounding) + ',' +
        d.target.x + ' ' + d.target.y + ',' + d.target.x;
  }

  graph.append('g')
      .attr('fill', 'none')
      .attr('stroke-width', 1.2)
      .selectAll('path')
      .data(tree_struct.links())
      .join('path')
      .attr('d', draw_single_edge)
      .attr(
          'stroke', d => (d.target === d.source.children[0]) ? '#0F0' : '#F00');
}
